# -*- coding: utf-8 -*-
import os
import psutil
import yaml
import shlex
import subprocess
import time
import win32serviceutil
import win32con
import win32api
import win32ui
import win32gui
import Logger

"""
此程序为WinServer平台下执行RPA命令，包括一下功能：
所需第三方模块:psutil、pywin32

def cmd_exec()  -->  使用subprocess.Popen执行cmd命令，阻塞获取返回结果

def get_process_info(process_name,process_pid)  -->  可通过pid或应用名称获取到程序信息,如pid、name、status信息

def start_win_process(process_path)  -->  开启一个windows进程程序

def close_win_process(process_name, process_pid)  -->  可根据process_name 或 process_pid 关闭一个进程

def restart_win_process(process_path)  -->  重启程序

def monitor_process_status()  -->  监控配置文件内程序运行状态

def windows_service_rpa(service_name,order)  -->  开始、关闭、重启 windows服务

def winShot(hwnd)  -->  输入程序句柄进行截图，将图片数据转为二进制格式进行返回

def list_port_usage(process_name)  -> 输入程序名称，查看程序端口占用情况
"""


class MessageStruct:
    pid = None
    name = None
    status = None
    started = None
    msg = None


def cmd_exec(cmd, start_cmd_window=False, timeout=10):
    """
    subprocess.Popen 方式执行cmd命令
    如果程序为cmd窗口程序，需要单独弹出一个cmd窗口，设置参数start_cmd_window=True
    """
    Logger.info('')
    result = list()
    Logger.info("execute terminal command:[%s] ,parameters start_cmd_window=[%s],timeout=[%s]" % (
        cmd, start_cmd_window, timeout))
    try:
        cmd_list = shlex.split(cmd)
        if start_cmd_window:
            subprocess.Popen(cmd,
                             creationflags=subprocess.CREATE_NEW_CONSOLE)
            time.sleep(1)
            Logger.info("Execute terminal command [%s] complete !" % cmd)
            return True
        proc = subprocess.Popen(cmd_list, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                universal_newlines=True)

    except Exception as e:
        Logger.error(e)
        Logger.error("Popen cmd_list failure,try to execute cmd add parameters cmd=True")
        return False
    try:
        proc.wait(timeout)
    except TimeoutError as e:
        proc.terminate()
        raise TimeoutError(str(e))
    else:
        out = proc.stdout.readlines()
        for line in out:
            # print(line, end='')
            result.append(line.strip())
        Logger.info("The terminal command [%s] execute result:" + "".join(result))
    return result


# 获取程序的基本信息
def get_process_info(process_name=None, process_pid=None, print_info=True):
    """
    可通过pid和应用名称获取到程序信息（pid、name、status）

    输入 pid ： 通过pid查找程序状态                    return pid、name、status
    输入name ： 通过name查找程序状态                   return pid、name、status
    输入pid、name ： 判断系统中 pid 和 name 是否对应上   return True or False

    pid -> int
    name -> str
    status ->str
    """
    Logger.info("Enter the query process info function,parameters is name=[%s],pid=[%s]" % (process_name, process_pid))
    if process_pid is not None and process_name is None:  # 通过pid 查找程序
        Logger.info("Query process info by pid")
        for proc in psutil.process_iter():
            if proc.pid == process_pid:
                message = MessageStruct()
                message.status = proc.status()
                message.pid = proc.pid
                message.name = proc.name()
                if print_info:
                    print(message.__dict__)
                Logger.info("Query result %s" % str(proc))
                return proc

    if process_name is not None and process_pid is None:  # 通过name查找程序，可能有多个程序
        Logger.info("Query process info by name")
        process_list = []
        for proc in psutil.process_iter():
            if proc.name() == process_name:
                process_list.append(proc)  # 将process对象存入process列表中
        for i in process_list[:1]:
            message = MessageStruct()
            message.status = '1'
            message.msg = i.status()
            if print_info:
                print(message.__dict__)

        return process_list

    if process_name and process_pid:
        result = {'status': psutil.pid_exists(process_pid)}
        return psutil.pid_exists(process_pid)


# 开启一个windows程序
def start_win_process(process_path, start_cmd_window=False, print_info=True):
    """
    0. 启动前先查询应用是否正在运行
    1. 启动 windows 程序
    2. 启动成功后，查看程序运行状态

    return True or False
    """

    # 解析程序名
    process_name = process_path.split('\\')[-1]
    Logger.info("Enter start win process function,parameters process_path=[%s],start_cmd_window=[%s]" % (
        process_path, start_cmd_window))

    # 检查当前程序是否正在运行
    result = get_process_info(process_name=process_name, print_info=False)
    if result.__len__() > 0:
        Logger.error("%s program is running" % process_name)  # 程序正在运行，需要做其他处理
        message = MessageStruct()
        message.status = '0'
        message.msg = "%s program is running" % process_name
        if print_info:
            print(message.__dict__)
        return False, result

    # 执行开启命令
    cmd_exec(process_path, start_cmd_window=start_cmd_window)
    # 通过程序名查看系统中程序运行状态
    process_info = get_process_info(process_name=process_name, print_info=False)  # 通过程序名查询可能会有多个程序
    if process_info.__len__() == 0:
        message = MessageStruct()
        message.status = '0'
        message.msg = "%s program starting failure" % process_name
        if print_info:
            print(message.__dict__)
        # print("%s program starting failure" % process_name)
        Logger.warn("%s program starting failure" % process_name)
        return False
    if process_info[0].name():
        message = MessageStruct()
        message.status = '1'
        message.msg = "%s program starting successful" % process_name
        if print_info:
            print(message.__dict__)
        Logger.info("%s program starting successful:" % process_name)
    return True, process_info  # 返回True 和 程序列表


# 关闭windows程序
def close_win_process(process_name=None, process_pid=None, print_info=True):
    """
    根据程序名来结束进程，结束所有进程
    根据程序pid结束进程，结束指定进程id

    杀死后，检查系统中知否存在程序
    :return True or False
    """
    Logger.info("Enter close win process function,parameters process_name=[%s],process_pid=[%s]" % (
        process_name, process_pid))
    if process_name:
        Logger.info("close win process by process_name")
        process_info = get_process_info(process_name=process_name, print_info=False)
        if len(process_info) == 0:
            message = MessageStruct()
            message.status = '0'
            message.msg = '%s not running' % process_name
            print(message.__dict__)
            exit()
        result = cmd_exec('taskkill /f /im %s' % process_name)
        time.sleep(0.01)
        process_info = get_process_info(process_name=process_name, print_info=False)
        if result.__len__() > 0 and process_info.__len__() == 0:
            Logger.info("%s close successful" % process_name)
            message = MessageStruct()
            message.status = '1'
            message.msg = '%s close successful' % process_name
            if print_info:
                print(message.__dict__)
            return True
        else:
            Logger.error("%s close fail" % process_name)
            message = MessageStruct()
            message.status = '0'
            message.msg = '%s close fail' % process_name
            if print_info:
                print(message.__dict__)
            return False

    if process_pid:
        Logger.info("close win process by process_pid")
        result = cmd_exec('taskkill /f /pid %s' % process_pid)
        time.sleep(0.05)
        if result.__len__() > 0 and not psutil.pid_exists(pid=process_pid):
            Logger.info("%s close successful" % process_pid)
            message = MessageStruct()
            message.status = '1'
            message.msg = '%s close successful' % process_pid
            if print_info:
                print(message.__dict__)
            return True
        else:
            Logger.error("%s close fail" % process_name)

            message = MessageStruct()
            message.status = '0'
            message.msg = '%s close fail' % process_name
            if print_info:
                print(message.__dict__)
    Logger.error("process close failure")
    return False


def restart_win_process(process_path=None, start_cmd_window=False):
    """
    重启程序
    1. 查看程序是否运行，如果程序未运行，提示程序运行状态有误
    2. 如果程序正在运行，直接关闭程序
    3. 启动程序
    4. 查看程序运行状态
    """

    Logger.info("Enter restart win process function,parameters process_path=[%s],start_cmd_window=[%s]" % (
        process_path, start_cmd_window))
    process_name = process_path.split('\\')[-1]  # 解析程序名
    process_list = get_process_info(process_name=process_name, print_info=False)  # 查看程序运行状态
    if len(process_list) == 0:  # 如果当前程序未运行，提示状态有误
        Logger.error('%s program not running' % process_name)
        message = MessageStruct()
        message.status = '0'
        message.msg = '%s program not running' % process_name
        print(message.__dict__)
        return False, 'program not running'
    if close_win_process(process_name=process_name, print_info=False):  # 如果关闭成功
        start_win_process(process_path=process_path, start_cmd_window=start_cmd_window, print_info=False)  # 开启程序
        result = get_process_info(process_name=process_name, print_info=False)  # 查看程序运行状态
        if len(result) == 0:
            Logger.error('%s program restart fail' % process_name)
            message = MessageStruct()
            message.status = '0'
            message.msg = '%s restart fail' % process_name
            print(message.__dict__)
            exit()

        Logger.error('%s program restart successful' % process_name)
        message = MessageStruct()
        message.status = '1'
        message.msg = '%s restart successful' % process_name
        print(message.__dict__)
        return True, 'program restart successful'


def monitor_process_status(print_info=True):
    """
    定时执行或循环执行监测程序
    监测指定程序运行状态，程序清单在config.yml文件中
    :return:
    """
    Logger.info("Enter mointor process status function")
    f = open('config.yml', 'r', encoding='utf8')
    yaml_result = f.read()
    yaml_data = yaml.load(yaml_result, Loader=yaml.FullLoader)
    # 获取所有需要监测的程序
    all_monitor_process = yaml_data['program_list']
    Logger.info("That process need to be monitor are " + " ".join(all_monitor_process))
    for key, value in all_monitor_process.items():
        process_name = value.split('\\')[-1]
        result = get_process_info(process_name=process_name, print_info=False)  # 查看当前程序运行状态
        if len(result) > 0 and result[0].status() == 'running':  # 如果存在程序，打印出程序基本信息
            Logger.info("%s process running，pid:[%s],name:[%s],status:[%s]" % (
                process_name, result[0].pid, result[0].name(), result[0].status()))
            message = MessageStruct()
            message.status = '1'
            message.msg = "%s process running，pid:[%s],name:[%s],status:[%s]" % (
                process_name, result[0].pid, result[0].name(), result[0].status())
            if print_info:
                print(message.__dict__)

            print("%s process running，pid:[%s],name:[%s],status:[%s]" % (
                process_name, result[0].pid, result[0].name(), result[0].status()))
        else:
            Logger.error("%s process not running" % process_name)
            message = MessageStruct()
            message.status = '0'
            message.msg = "%s process not running" % process_name
            if print_info:
                print(message.__dict__)


def windows_service_rpa(server_name, order):
    """
    启动、关闭、重启 windows服务
    server_name ： 服务名称
    order ： 输入命令 start restart stop -> str
    """
    Logger.info("Entry windows service rpa function,parameters server_name=[%s],order=[%s]" % (server_name, order))
    if order == 'start':
        try:
            Logger.info("start %s" % server_name)
            win32serviceutil.StartService(server_name)
            # 查看服务运行状态
            result = cmd_exec('sc query %s' % server_name)
            message = MessageStruct()
            message.status = '1'
            message.msg = 'service:[%s] start successful' % server_name
            print(message.__dict__)
            return True
        except Exception as e:
            message = MessageStruct()
            message.status = '0'
            message.msg = e.args
            print(message.__dict__)
            Logger.error(e.args)
            return False, e.args
    if order == 'restart':
        try:
            Logger.info("restart %s" % server_name)

            win32serviceutil.RestartService(server_name)
            cmd_exec('sc query %s' % server_name)
            message = MessageStruct()
            message.status = '1'
            message.msg = 'service:[%s] restart successful' % server_name
            print(message.__dict__)
            return True, 'service:[%s] restart successful' % server_name
        except Exception as e:
            message = MessageStruct()
            message.status = '0'
            message.msg = e.args
            print(message.__dict__)
            Logger.error(e.args)
            return False, e.args

    if order == 'stop':
        try:
            Logger.info("stop %s" % server_name)

            win32serviceutil.StopService(server_name)
            cmd_exec('sc query %s' % server_name)
            message = MessageStruct()
            message.status = '1'
            message.msg = 'service:[%s] stop successful' % server_name
            print(message.__dict__)
            return True, 'service:[%s] stop successful' % server_name
        except Exception as e:
            message = MessageStruct()
            message.status = '0'
            message.msg = e.args
            print(message.__dict__)
            Logger.error(e.args)
            return False, e.args
    return False, 'Incorrect order input'


def winShot(process_name, file_path):
    """
     将窗口置顶
     根据窗口句柄截取窗口视图
     返回图片二进制数据
    """
    Logger.info("Entry process screenshot function，parameters hwnd=[%s],file_path=[%s]" % (process_name, file_path))
    hwnd = win32gui.FindWindow(0, process_name)
    try:
        win32gui.SetForegroundWindow(hwnd)
    except Exception as e:
        message = MessageStruct()
        message.status = '0'
        message.msg = "当前应用不存在，请再次确认proces_name:[%s]输入是否正确或程序在后台运行，没有窗口" % process_name
        print(message.__dict__)
        Logger.error(e)
        exit()
    time.sleep(0.01)
    win32gui.ShowWindow(hwnd, win32con.SW_SHOWNORMAL)
    r = win32gui.GetWindowRect(hwnd)
    hwin = win32gui.GetDesktopWindow()
    # 图片最左边距离主屏左上角的水平距离
    left = win32api.GetSystemMetrics(win32con.SM_XVIRTUALSCREEN)
    # 图片最上边距离主屏左上角的垂直距离
    top = win32api.GetSystemMetrics(win32con.SM_YVIRTUALSCREEN)
    hwindc = win32gui.GetWindowDC(hwin)
    srcdc = win32ui.CreateDCFromHandle(hwindc)
    memdc = srcdc.CreateCompatibleDC()
    bmp = win32ui.CreateBitmap()
    bmp.CreateCompatibleBitmap(srcdc, r[2] - r[0], r[3] - r[1])
    memdc.SelectObject(bmp)
    memdc.BitBlt((-r[0], top - r[1]), (r[2], r[3] - top), srcdc, (left, top), win32con.SRCCOPY)
    bmp.SaveBitmapFile(memdc, file_path)
    if os.path.getsize(file_path) > 20480:
        message = MessageStruct()
        message.status = '1'
        message.msg = "%s 程序快照获取成功，文件保存路径：%s" % (process_name, file_path)
        print(message.__dict__)
        # print("[%s] 程序快照获取成功，文件保存路径：[%s]" % (process_name, pic_save_path))
        Logger.info("[%s] 程序快照获取成功，文件保存路径：%s" % (process_name, file_path))
    else:
        print("程序快照获取失败")
        os.remove(file_path)
        message = MessageStruct()
        message.status = '0'
        message.msg = '程序快照获取失败'
        print(message.__dict__)
        Logger.error("程序快照获取失败")


def list_port_usage(process_name):
    Logger.info("Entry list port usage function,parameters process_name=[%s]" % process_name)
    """
    查看windos程序端口占用情况
    process_name 程序名称
    """
    # 1. 通过程序名称，获取程序pid
    ports = []
    process_infos = get_process_info(process_name=process_name, print_info=False)

    for process in process_infos:
        pid = process.pid
        port = cmd_exec('netstat -aon|findstr "%s"' % pid)
        for i in port:
            ports.append(i.split('    ')[1])
    result = list(set(ports))
    message = MessageStruct()
    message.status = '1'
    message.msg = "The %s process is running port:" % process_name + " ".join(result)
    print(message.__dict__)
    Logger.info("The %s process is running port:" % process_name + " ".join(result))


if __name__ == '__main__':
    import argparse

    funcs = ['get_process_info', 'start_win_process', 'close_win_process', 'restart_win_process',
             'monitor_process_status', 'windows_service_rpa', 'winShot']

    parser = argparse.ArgumentParser(
        description='This is a package for executing RPA commands，It includes the following functions')
    parser.add_argument('--func', default=None, help='executing function name')
    parser.add_argument('--process_name', default=None, help='executing process name')
    parser.add_argument('--process_pid', default=None, help='executing process pid')
    parser.add_argument('--process_path', default=None, help='executing process path')
    parser.add_argument('--service_name', default=None, help='windows server name')
    parser.add_argument('--service_order', default=None, help='executing windows server order:[start,close,restart]')
    parser.add_argument('--start_cmd_window', default=False, help='only use starting a program,input True or False')
    parser.add_argument('--pic_save_path', default=False, help='Path for saving snapshots of applications')
    args = parser.parse_args()
    Logger.info("process start...")
    func = args.func
    process_name = args.process_name
    process_pid = args.process_pid
    process_path = args.process_path
    service_name = args.service_name
    service_order = args.service_order
    start_cmd_window = args.start_cmd_window
    pic_save_path = args.pic_save_path

    if func not in funcs:
        Logger.error("func parameters error,this process support function is ：\n" + "\n".join(funcs))
        print("func parameters error,this process support function is ：\n" + "\n".join(funcs))

    if func == 'get_process_info':
        get_process_info(process_name=process_name, process_pid=process_pid)

    if func == 'start_win_process':
        if start_cmd_window == 'True':
            start_win_process(process_path=process_path, start_cmd_window=True)
        if start_cmd_window == 'False':
            start_win_process(process_path=process_path, start_cmd_window=False)
    if func == 'close_win_process':
        close_win_process(process_name=process_name, process_pid=process_pid)

    if func == 'restart_win_process':
        if start_cmd_window == 'True':
            restart_win_process(process_path=process_path, start_cmd_window=True)
        if start_cmd_window == 'False':
            restart_win_process(process_path=process_path, start_cmd_window=False)

    if func == 'monitor_process_status':
        monitor_process_status()

    if func == 'windows_service_rpa':
        windows_service_rpa(server_name=service_name, order=service_order)

    if func == 'winShot':
        winShot(process_name=process_name, file_path=pic_save_path)

    if func == 'list_port_usage':
        list_port_usage(process_name=process_name)
